# -*- encoding: utf-8 -*-
"""
@File   : SuperMapPerformance.py
@Author : @li-xin12345
@Date   : 2022/08/16
"""
import os
import sys
import time
import subprocess
import shutil


class SuperMapPerformance:
    def __init__(self, res_dir_path, host, port, username, env):
        self.res_dir_path = res_dir_path
        self.host = host
        self.port = port
        self.username = username
        self.env_path = env

    def create_table(self, test_db):
        create_cmd = f'begin;' \
            f'create table t_polygon_gis (id int primary key, pgis geometry);' \
            f'with a as (select id, (random()*10000)::int x, ' \
            f'(random()*10000)::int y from ' \
            f'generate_series(1,100000000) t(id)) ' \
            f'insert into t_polygon_gis ' \
            f'select id, st_makeenvelope(x, y, x+1, y+1) from a;' \
            f'end;' \
            f'create Index t_polygon_gis_idx_geom ' \
            f'on t_polygon_gis using GiST(pgis);'
        execute_cmd = f"source {self.env_path};" \
            f"gsql -d {test_db} -p {self.port} -r -c '{create_cmd}'"
        execute_msg = subprocess.getoutput(execute_cmd)
        return execute_msg

    def create_function(self, test_db, create_function_file):
        execute_cmd = f"source {self.env_path};" \
            f"gsql -d {test_db} -p {self.port} -r -f {create_function_file}"
        execute_msg = subprocess.getoutput(execute_cmd)
        return execute_msg

    def run_pgbench(self, tool_path, db, pwd, filename,
                    duration='20', clients='1', threads='1'):
        res = '_'.join([filename.split('.')[0], 'result.log'])
        res_path = os.path.join(self.res_dir_path, res)
        cmd = f'rm -rf {res_path};touch {res_path}'
        os.system(cmd)
        # 执行测试脚本
        run_cmd = f'{tool_path} {db} -r -n -T {duration} -c {clients} ' \
            f'-j {threads} -h {self.host} -p {self.port} ' \
            f'-U {self.username} -P 5 -f {filename}'
        execute_cmd = f'''expect <<EOF
                        set timeout 120
                        spawn {run_cmd}
                        expect "Password:"
                        send "{pwd}\\n"
                        expect eof\n''' + '''EOF'''
        execute_msg = subprocess.getoutput(execute_cmd)
        with open(res_path, 'a+') as f:
            f.write(execute_msg)
            f.close()
        return res_path

    def read_result(self, result_file):
        # 读取执行结果
        ls_cmd = f'ls {self.res_dir_path}'
        log_results = subprocess.getoutput(ls_cmd).splitlines()
        for log_result in log_results:
            log_result_path = os.path.join(self.res_dir_path, log_result)
            tps_cmd = ['grep', 'excluding', log_result_path]
            process_2 = subprocess.Popen(tps_cmd, shell=False,
                                         stdout=subprocess.PIPE,
                                         stderr=subprocess.PIPE)
            tps = process_2.stdout.read().decode('utf-8').strip()
            process_2.stdout.close()
            errors = process_2.stderr.read().decode('utf-8')
            process_2.stderr.close()
            # 将tps值并写入文件
            with open(result_file, 'a+') as file2:
                if errors:
                    file2.write("{0}:{1}\n".format(log_result, errors))
                else:
                    file2.write("{0}:{1}\n".format(log_result, tps.split()[2]))
                file2.close()
        return result_file


if __name__ == '__main__':
    pgbench_path = sys.argv[1]
    db_name = sys.argv[2]
    db_host = sys.argv[3]
    db_port = sys.argv[4]
    db_user = sys.argv[5]
    env_path = sys.argv[6]
    passwd = sys.argv[7]
    func_list = ['create_function_1.sql', 'create_function_2.sql',
                 'create_function_3.sql', 'create_function_4.sql']
    # 创建结果保存文件
    date_res = time.strftime("%Y-%m-%d", time.localtime())
    dir_name = '_'.join(['dir_result', date_res])
    dir_path = os.path.join(os.getcwd(), dir_name)
    if os.path.exists(dir_path):
        shutil.rmtree(dir_path)
        os.makedirs(dir_path)
    else:
        os.makedirs(dir_path)
    res_file_name = ''.join(['res_file_', date_res, '.txt'])
    touch_cmd = f'rm -f {res_file_name};touch {res_file_name};'
    subprocess.getoutput(touch_cmd)
    test = SuperMapPerformance(dir_path, db_host, db_port, db_user, env_path)
    # test1:1亿个box和随机点相交测试
    exec_cmd = f"source {env_path};gsql -d {db_name} -p {db_port} -r -c '\\d'"
    exec_msg = subprocess.getoutput(exec_cmd)
    if 't_polygon_gis' not in exec_msg:
        create_result = test.create_table(db_name)
        if 'CREATE TABLE' and 'CREATE INDEX' in create_result:
            test.run_pgbench(pgbench_path, db_name, passwd, 'test1.sql')
        else:
            print('创建表失败' + create_result)
    else:
        test.run_pgbench(pgbench_path, db_name, passwd, 'test1.sql')
    # test2_1:面图层框选查询
    create_func_msg_1 = test.create_function(db_name, func_list[0])
    if 'CREATE FUNCTION' in create_func_msg_1:
            shutil.copy('test2.sql', 'test2_1.sql')
            test.run_pgbench(pgbench_path, db_name, passwd, 'test2_1.sql')
    else:
        print(create_func_msg_1)
    create_func_msg_2 = test.create_function(db_name, func_list[1])
    if 'CREATE FUNCTION' in create_func_msg_2:
        shutil.copy('test2.sql', 'test2_2.sql')
        test.run_pgbench(pgbench_path, db_name, passwd, 'test2_2.sql')
    else:
        print(create_func_msg_2)
    # test2_2:线图层框选查询
    create_func_msg_3 = test.create_function(db_name, func_list[2])
    if 'CREATE FUNCTION' in create_func_msg_3:
        shutil.copy('test2.sql', 'test2_3.sql')
        test.run_pgbench(pgbench_path, db_name, passwd, 'test2_3.sql')
    else:
        print(create_func_msg_3)
    create_func_msg_4 = test.create_function(db_name, func_list[3])
    if 'CREATE FUNCTION' in create_func_msg_4:
        shutil.copy('test2.sql', 'test2_4.sql')
        test.run_pgbench(pgbench_path, db_name, passwd, 'test2_4.sql')
    else:
        print(create_func_msg_4)
    # test3:点面叠加-包含
    test.run_pgbench(pgbench_path, db_name, passwd, 'test3.sql')
    # test4:面面叠加-裁剪
    test.run_pgbench(pgbench_path, db_name, passwd, 'test4.sql')
    # test5_1:缓冲区分析10000
    test.run_pgbench(pgbench_path, db_name, passwd, 'test5_1.sql')
    # test5_2:缓冲区分析100000
    test.run_pgbench(pgbench_path, db_name, passwd, 'test5_2.sql')
    # test6:属性查询
    test.run_pgbench(pgbench_path, db_name, passwd, 'test6.sql')
    # test7:主键查询
    test.run_pgbench(pgbench_path, db_name, passwd, 'test7.sql')
    # 读取所有结果并写入文件中
    test.read_result(res_file_name)
